package org.finesys.chat.core.base.embedding.store.pgvector;

import java.sql.Array;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.stream.Collectors;

import javax.sql.DataSource;

import org.finesys.chat.core.base.document.Metadata;
import org.finesys.chat.core.base.embedding.Embedding;
import org.finesys.chat.core.base.embedding.EmbeddingMatch;
import org.finesys.chat.core.base.embedding.EmbeddingSearchRequest;
import org.finesys.chat.core.base.embedding.EmbeddingSearchResult;
import org.finesys.chat.core.base.embedding.EmbeddingStore;
import org.finesys.chat.core.base.embedding.filter.Filter;
import org.finesys.chat.core.base.embedding.store.pgvector.constants.IVFFlatIndexEnum;
import org.finesys.chat.core.base.embedding.store.pgvector.factory.MetadataHandlerFactory;
import org.finesys.chat.core.base.embedding.store.util.Generator;
import org.finesys.chat.core.base.segment.TextSegment;
import org.finesys.common.core.util.Utils;
import org.finesys.common.core.util.ValidationUtil;
import org.postgresql.ds.PGSimpleDataSource;

import com.pgvector.PGvector;

public class PgVectorEmbeddingStore implements EmbeddingStore<TextSegment>, AutoCloseable {

    private final DataSource datasource;
    private final String table;
    private final MetadataHandler metadataHandler;


    public PgVectorEmbeddingStore(DataSource datasource, String table, Integer dimension, Boolean useIndex, Integer indexListSize, Boolean createTable, Boolean dropTableFirst, MetadataStorageConfig metadataStorageConfig) {
        this.datasource = ValidationUtil.ensureNotNull(datasource, "datasource");
        this.table = ValidationUtil.ensureNotBlank(table, "table");
        MetadataStorageConfig config = Utils.getOrDefault(metadataStorageConfig, DefaultMetadataStorageConfig.defaultConfig());
        this.metadataHandler = MetadataHandlerFactory.get(config);
        useIndex = Utils.getOrDefault(useIndex, false);
        createTable = Utils.getOrDefault(createTable, true);
        dropTableFirst = Utils.getOrDefault(dropTableFirst, false);
        this.initTable(dropTableFirst, createTable, useIndex, dimension, indexListSize);
    }

    public PgVectorEmbeddingStore(String host, int port, String userName, String password, String database, String table, Integer dimension, Boolean useIndex, Integer indexListSize, Boolean createTable, Boolean dropTableFirst, MetadataStorageConfig metadataStorageConfig) {
        this(createDataSource(host, port, userName, password, database), table, dimension, useIndex, indexListSize, createTable, dropTableFirst, metadataStorageConfig);
    }

    private static DataSource createDataSource(String host, int port, String userName, String password, String database) {
        host = ValidationUtil.ensureNotBlank(host, "host");
        port = ValidationUtil.ensureGreaterThanZero(port, "port");
        userName = ValidationUtil.ensureNotBlank(userName, "userName");
        password = ValidationUtil.ensureNotBlank(password, "password");
        database = ValidationUtil.ensureNotBlank(database, "database");
        PGSimpleDataSource pgSimpleDataSource = new PGSimpleDataSource();
        pgSimpleDataSource.setServerNames(new String[]{host});
        pgSimpleDataSource.setPortNumbers(new int[]{port});
        pgSimpleDataSource.setDatabaseName(database);
        pgSimpleDataSource.setUser(userName);
        pgSimpleDataSource.setPassword(password);
        return pgSimpleDataSource;
    }

    private void initTable(Boolean dropTableFirst, Boolean createTable, Boolean useIndex, Integer dimension, Integer indexListSize) {
        String query = "init";
        Connection connection = null;
        Statement statement = null;
        try {
            connection = this.getConnection();
            statement = connection.createStatement();
            if (dropTableFirst) {
                statement.executeUpdate("DROP TABLE IF EXISTS " + this.table);
            }
            // 创建表，包含一个向量列和一个文本内容列。此处可以根据实际情况调整字段类型和名称。
            //CREATE TABLE items (id bigserial PRIMARY KEY, embedding vector(3));
            if (createTable) {
                query = String.format("CREATE TABLE IF NOT EXISTS %s (id UUID PRIMARY KEY, embedding VECTOR(%d),text TEXT, %s);",
                        this.table, ValidationUtil.ensureGreaterThanZero(dimension, "dimension"),
                        metadataHandler.columnDefinitionsString());
                statement.executeUpdate(query);
            }
            // 创建索引，使用IVF-Flat算法进行近似最近邻搜索。此处可以根据实际情况调整参数。
            // 创建索引，使用HNSW算法进行近似最近邻搜索。此处可以根据实际情况调整参数。
            //CREATE INDEX ON items USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
            if (useIndex) {
                String indexName = String.format("%s_ivfflat_index", this.table);
                query = String.format("CREATE INDEX IF NOT EXISTS %s ON %s USING ivfflat(embedding %s) WITH (lists = %s)"
                        , indexName, this.table, IVFFlatIndexEnum.COSINE_DISTANCE.getValue(), ValidationUtil.ensureGreaterThanZero(indexListSize, "indexListSize"));
                statement.executeUpdate(query);
            }
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            //关闭Statement对象
            try {
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            //关闭连接
            try {
                if (connection != null) {
                    connection.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }
    //Update vectors
    //INSERT INTO items (id, embedding) VALUES (1, '[1,2,3]'), (2, '[4,5,6]')
    //    ON CONFLICT (id) DO UPDATE SET embedding = EXCLUDED.embedding;

    //Supported distance functions are:
    //<-> - L2 distance
    //<#> - (negative) inner product
    //<=> - cosine distance
    //<+> - L1 distance
    //<~> - Hamming distance (binary vectors)
    //<%> - Jaccard distance (binary vectors)

    private Connection getConnection() throws SQLException {
        Connection connection = null;
        Statement statement = null;
        try {
            connection = this.datasource.getConnection();
            statement = connection.createStatement();
            statement.executeUpdate("CREATE EXTENSION IF NOT EXISTS vector ");
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            if (statement != null) {
                statement.close();
            }
        }
        PGvector.addVectorType(connection);
        return connection;
    }

    @Override
    public String add(Embedding embedding) {
        String id = UUID.randomUUID().toString();
        this.addInternal(id, embedding, (TextSegment) null);
        return id;
    }

    @Override
    public String add(Embedding embedding, String key) {
        this.addInternal(key, embedding, (TextSegment) null);
        return key;
    }

    @Override
    public String add(Embedding embedding, TextSegment textSegment) {
        String id = UUID.randomUUID().toString();
        this.addInternal(id, embedding, textSegment);
        return id;
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings) {
        List<String> ids = embeddings.stream().map(embedding -> {
            String id = UUID.randomUUID().toString();
            this.addInternal(id, embedding, (TextSegment) null);
            return id;
        }).collect(Collectors.toList());
        return ids;
    }

    @Override
    public List<String> addAll(List<Embedding> embeddings, List<TextSegment> textSegments) {
        List<String> ids = Generator.generateRandomIds(embeddings.size());
        return this.addAll(ids, embeddings, textSegments);
    }

    @Override
    public List<String> addAll(List<String> ids, List<Embedding> embeddings, List<TextSegment> textSegments) {

        if (!Utils.isNullOrEmpty(ids) && !Utils.isNullOrEmpty(embeddings)) {
            ValidationUtil.ensureTrue(ids.size() == embeddings.size(), "ids size is not equal to embeddings size");
            ValidationUtil.ensureTrue(textSegments == null || embeddings.size() == textSegments.size(), "embeddings size is not equal to embedded size");

            Connection connection = null;
            PreparedStatement statement = null;
            //sql
            String insertSql = String.format("INSERT INTO %s (id, embedding,text,%s) VALUES (?, ?,?,%s) " +
                            "on conflict(id) do update set embedding=EXCLUDED.embedding,text=EXCLUDED.text,%s;",
                    this.table, String.join(",", metadataHandler.columnsNames())
                    , String.join(",", Collections.nCopies(this.metadataHandler.columnsNames().size(), "?"))
                    , this.metadataHandler.insertClause());

            try {
                connection = this.getConnection();
                statement = connection.prepareStatement(insertSql);

                int i = 0;
                while (true) {
                    if (i >= ids.size()) {
                        statement.executeBatch();
                    }
                    statement.setObject(1, UUID.fromString(ids.get(i)));
                    statement.setObject(2, new PGvector(embeddings.get(i).getVectorFloat()));
                    if (textSegments != null && textSegments.get(i) != null) {
                        statement.setObject(3, textSegments.get(i).getText());
                        this.metadataHandler.setMetadata(statement, 4, textSegments.get(i).getMetadata());
                    } else {
                        statement.setNull(3, 12);
                        for (int j = 4; j < this.metadataHandler.columnsNames().size(); j++) {
                            statement.setNull(j, 12);
                        }
                    }
                    statement.addBatch();
                    i++;
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            } finally {
                //关闭Statement对象
                try {
                    if (statement != null) {
                        statement.close();
                    }
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
                //关闭连接
                try {
                    if (connection != null) {
                        connection.close();
                    }
                } catch (SQLException e) {
                    throw new RuntimeException(e);
                }
            }
        }

        return ids;
    }

    private void addInternal(String id, Embedding embedding, TextSegment textSegment) {
        this.addAll(Collections.singletonList(id), Collections.singletonList(embedding), textSegment == null ? null : Collections.singletonList(textSegment));
    }

    @Override
    public void delete(String key) {
        ValidationUtil.ensureNotBlank(key, "key");
        String sql = String.format("DELETE FROM %s WHERE id = ?", this.table);
        Connection connection = null;
        PreparedStatement statement = null;
        try {
            connection = this.getConnection();
            statement = connection.prepareStatement(sql);
            statement.setString(1, key);
            statement.executeUpdate();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            //关闭Statement对象
            try {
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            //关闭连接
            try {
                if (connection != null) {
                    connection.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public void deleteAll(List<String> keys) {
        ValidationUtil.ensureNotEmpty(keys, "keys");
        String sql = String.format("DELETE FROM %s WHERE id = Any(?)", this.table);
        Connection connection = null;
        PreparedStatement statement = null;
        try {
            connection = this.getConnection();
            statement = connection.prepareStatement(sql);
            Array array = connection.createArrayOf("uuid", keys.stream().map(UUID::fromString).toArray());
            statement.setArray(1, array);
            statement.executeUpdate();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            //关闭Statement对象
            try {
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            //关闭连接
            try {
                if (connection != null) {
                    connection.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public void deleteAll(Filter filter) {
        ValidationUtil.ensureNotNull(filter, "filter");
        String whereClause = this.metadataHandler.whereClause(filter);
        String sql = String.format("DELETE FROM %s WHERE %s", this.table, whereClause);
        Connection connection = null;
        PreparedStatement statement = null;
        try {
            connection = this.getConnection();
            statement = connection.prepareStatement(sql);
            statement.executeUpdate();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            //关闭Statement对象
            try {
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            //关闭连接
            try {
                if (connection != null) {
                    connection.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public void deleteAll() {
        String sql = String.format("TRUNCATE TABLE %s ", this.table);
        Connection connection = null;
        PreparedStatement statement = null;
        try {
            connection = this.getConnection();
            statement = connection.prepareStatement(sql);
            statement.executeUpdate();
        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            //关闭Statement对象
            try {
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            //关闭连接
            try {
                if (connection != null) {
                    connection.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public EmbeddingSearchResult<TextSegment> search(EmbeddingSearchRequest request) {
        Embedding embedding = request.getQueryEmbedding();
        int maxLimit = request.getLimit();
        double minScore = request.getMinScore();
        Filter filter = request.getFilter();

        String referenceVector = Arrays.toString(embedding.getVector());
        String whereClause = filter == null ? "" : this.metadataHandler.whereClause(filter);
        whereClause = whereClause.isEmpty() ? "" : "AND " + whereClause;
        //组装查询语句
        String query = String.format("SELECT (2 - (embedding <=> '%s')) / 2 AS score, id, embedding, text, %s FROM %s " +
                        "WHERE round(cast(float8 (embedding <=> '%s') as numeric), 8) <= round(2 - 2 * %s, 8) %s " +
                        "ORDER BY embedding <=> '%s' LIMIT %d;", referenceVector, String.join(",", this.metadataHandler.columnsNames()),
                this.table, referenceVector, minScore, whereClause, referenceVector, maxLimit);

        Connection connection = null;
        PreparedStatement statement = null;
        ResultSet resultSet = null;
        //
        List<EmbeddingMatch<TextSegment>> result = new ArrayList<>();
        try {
            connection = this.getConnection();
            statement = connection.prepareStatement(query);
            resultSet = statement.executeQuery();
            Integer index = 0;
            while (resultSet != null && resultSet.next()) {
                String id = resultSet.getString("id");
                Double score = resultSet.getDouble("score");
                PGvector vector = resultSet.getObject("embedding", PGvector.class);
                Embedding embedding1 = Embedding.from(vector.toArray());
                String text = resultSet.getString("text");
                TextSegment textSegment = null;
                if (Utils.isNotNullOrBlank(text)) {
                    Metadata metadata = this.metadataHandler.fromResultSet(resultSet);
                    textSegment = TextSegment.from(text, metadata);
                }
                result.add(new EmbeddingMatch<>(score, id, embedding1, textSegment, index));
                index++;
            }

        } catch (SQLException e) {
            throw new RuntimeException(e);
        } finally {
            //关闭ResultSet对象
            try {
                if (resultSet != null) {
                    resultSet.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            //关闭Statement对象
            try {
                if (statement != null) {
                    statement.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
            //关闭连接
            try {
                if (connection != null) {
                    connection.close();
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
        return new EmbeddingSearchResult<>(result);
    }


    public static DatasourceBuilder datasourceBuilder() {
        return new DatasourceBuilder();
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class DatasourceBuilder {
        private DataSource dataSource;
        private String table;
        private Integer dimension;
        private Boolean useIndex;
        private Integer indexListSize;
        private Boolean createTable;
        private Boolean dropTableFirst;
        private MetadataStorageConfig metadataStorageConfig;

        public DatasourceBuilder() {
        }

        public DatasourceBuilder dataSource(DataSource val) {
            this.dataSource = val;
            return this;
        }

        public DatasourceBuilder table(String val) {
            this.table = val;
            return this;
        }

        public DatasourceBuilder dimension(Integer val) {
            this.dimension = val;
            return this;
        }

        public DatasourceBuilder useIndex(Boolean val) {
            this.useIndex = val;
            return this;
        }

        public DatasourceBuilder indexListSize(Integer val) {
            this.indexListSize = val;
            return this;
        }

        public DatasourceBuilder createTable(Boolean val) {
            this.createTable = val;
            return this;
        }

        public DatasourceBuilder dropTableFirst(Boolean val) {
            this.dropTableFirst = val;
            return this;
        }

        public DatasourceBuilder metadataStorageConfig(MetadataStorageConfig val) {
            this.metadataStorageConfig = val;
            return this;
        }

        public PgVectorEmbeddingStore build() {
            return new PgVectorEmbeddingStore(this.dataSource, this.table, this.dimension, this.useIndex, this.indexListSize, this.createTable, this.dropTableFirst, this.metadataStorageConfig);
        }
    }

    public static class Builder {

        private String host;

        private Integer port;

        private String user;

        private String password;

        private String database;

        private String table;

        private Integer dimension;

        private Boolean useIndex;

        private Integer indexListSize;

        private Boolean createTable;

        private Boolean dropTableFirst;

        private MetadataStorageConfig metadataStorageConfig;

        Builder() {

        }

        public Builder host(String host) {
            this.host = host;
            return this;
        }

        public Builder port(Integer port) {
            this.port = port;
            return this;
        }

        public Builder user(String user) {
            this.user = user;
            return this;
        }

        public Builder password(String password) {
            this.password = password;
            return this;
        }

        public Builder database(String database) {
            this.database = database;
            return this;
        }

        public Builder table(String table) {
            this.table = table;
            return this;
        }

        public Builder dimension(Integer dimension) {
            this.dimension = dimension;
            return this;
        }

        public Builder useIndex(Boolean useIndex) {
            this.useIndex = useIndex;
            return this;
        }

        public Builder indexListSize(Integer indexListSize) {
            this.indexListSize = indexListSize;
            return this;
        }

        public Builder createTable(Boolean createTable) {
            this.createTable = createTable;
            return this;
        }

        public Builder dropTableFirst(Boolean dropTableFirst) {
            this.dropTableFirst = dropTableFirst;
            return this;
        }

        public Builder metadataStorageConfig(MetadataStorageConfig metadataStorageConfig) {
            this.metadataStorageConfig = metadataStorageConfig;
            return this;
        }

        public PgVectorEmbeddingStore build() {
            return new PgVectorEmbeddingStore(this.host, this.port, this.user, this.password, this.database, this.table, this.dimension, this.useIndex, this.indexListSize, this.createTable, this.dropTableFirst, this.metadataStorageConfig);
        }
    }


    @Override
    public void close() throws Exception {
        if (this.datasource != null && this.datasource.getConnection() != null) {
            this.datasource.getConnection().close();
        }
    }

    @Override
    public void stop() {
        try {
            this.close();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
